#pragma once
#include <vector>
#include "MyUnit.hpp"
#include "defines.hpp"
#include "Line.hpp"
#include "BattlePredictor.hpp"

template <class T>
class Squad{
protected:
	std::set<T*> m_units;
	int gotoFormation(
		const std::vector<BWAPI::Position> &posVector
	);
public:
	Squad();
	void addUnit(
		BWAPI::Unit* unit
	);
	void removeUnit(
		BWAPI::Unit* unit
	);
	void onFrame();
	BWAPI::Position getPosition() const;
	virtual void setBackPointer(T* unit)=0;
	virtual void onFrame_() {};
	virtual void drawStats(
		const int x,
		const int y
	) {};
	int size() const { return m_units.size(); }
	int formLine(
		const Line &l
	);
	void addToVector(
		std::vector<BWAPI::Unit*> &v
	) const;
	void addToBattlePredictor(
		BattlePredictor &battlePredictor
	) const;
};

template <class T>
Squad<T>::Squad()
{
}

template <class T>
void Squad<T>::addUnit(
	BWAPI::Unit* unit
)
{
	if (unit->getType() == UnitTypes::Special_Map_Revealer) {
		return;
	}
	T *u = new T(unit);
	setBackPointer(u);
	m_units.insert( u );
}

template <class T>
void Squad<T>::removeUnit(
	BWAPI::Unit* unit
)
{
	for (std::set<T*>::iterator i=m_units.begin() ; i!=m_units.end() ; i++) {
		if ((*i)->getBWAPIUnit() == unit) {
			m_units.erase(i);
			break;
		}
	}
}

template <class T>
void Squad<T>::onFrame()
{
	onFrame_();
#ifdef DISPLAY_UNIT_COMMAND
	for (std::set<T*>::iterator i=m_units.begin() ; i!=m_units.end() ; i++) {
		(*i)->displayLastCommand();
	}
#endif
}

template <class T>
BWAPI::Position Squad<T>::getPosition() const
{
	BWAPI::Position p(0,0);
	int counter = 0;
	for (std::set<T*>::const_iterator i=m_units.begin() ; i!=m_units.end() ; i++) {
		p += (*i)->getBWAPIUnit()->getPosition();
		counter++;
	}
	p.x() /= double(counter);
	p.y() /= double(counter);
	return p;
}

template <class T>
int Squad<T>::gotoFormation(
	const std::vector<BWAPI::Position> &posVector
)
{
	if (posVector.size() != m_units.size()) {
		Broodwar->printf("ERROR in gotoFormation(): posVector.size() != m_units.size()");
		return 0;
	}
	std::vector< std::vector<int> > dMatrix(m_units.size(), std::vector<int>(m_units.size(), 100000));
	int idx=0;
	for (std::set<T*>::iterator i=m_units.begin() ; i!=m_units.end() ; i++, idx++) {
		BWAPI::Position p = (*i)->getBWAPIUnit()->getPosition();
		for (int j=0 ; j<posVector.size() ; j++) {
			dMatrix[idx][j] = CALC_D2(p,posVector[j]);
		}
	}
	int sumD2 = 0;
	std::vector<bool> availableUnits(m_units.size(), true);
	std::vector<bool> availablePos(m_units.size(), true);
	for (int k=0 ; k<posVector.size() ; k++) {
		int maxMinD2 = -1;
		T* maxMinDistanceUnit = NULL;
		int maxMinDistanceUnitIdx = 0;
		int maxMinDistancePosIdx = 0;
		for (int j=0 ; j<posVector.size() ; j++) {
			if (!availablePos[j]) {
				continue;
			}
			int minDistance = 100000;
			T* minDistanceUnit = NULL;
			int minDistanceUnitIdx = 0;
			idx = 0;
			for (std::set<T*>::iterator i=m_units.begin() ; i!=m_units.end() ; i++, idx++) {
				if (!availableUnits[idx]) {
					continue;
				}
				if (dMatrix[idx][j]<=minDistance) {
					minDistance = dMatrix[idx][j];
					minDistanceUnit = *i;
					minDistanceUnitIdx = idx;
				}
			}
			if (minDistance>maxMinD2) {
				maxMinD2 = minDistance;
				maxMinDistanceUnit = minDistanceUnit;
				maxMinDistancePosIdx = j;
				maxMinDistanceUnitIdx = minDistanceUnitIdx;
			}
		}
		maxMinDistanceUnit->getBWAPIUnit()->move( posVector[maxMinDistancePosIdx] );
		availablePos[maxMinDistancePosIdx] = false;
		availableUnits[maxMinDistanceUnitIdx] = false;
		sumD2 += maxMinD2;
	}
	return (sumD2 / m_units.size());
}

template <class T>
int Squad<T>::formLine(
		const Line &l
)
{
	if (m_units.size()==0) {
		return 0;
	}
	BWAPI::Position p1, p2;
	l.endPoints(p1, p2);
	int diffX = p2.x()-p1.x();
	int diffY = p2.y()-p1.y();
	std::vector<BWAPI::Position> posVector(m_units.size());
	for (int i=0 ; i<m_units.size() ; i++) {
		posVector[i] = BWAPI::Position(int(p1.x())+diffX*i/int(m_units.size()-1), int(p1.y())+diffY*i/int(m_units.size()-1));
	}
	return gotoFormation( posVector );
}

template <class T>
void Squad<T>::addToVector(
	std::vector<BWAPI::Unit*> &v
) const
{
	for (std::set<T*>::const_iterator i=m_units.begin() ; i!=m_units.end() ; i++) {
		v.push_back( (*i)->getBWAPIUnit() );
	}
}

template <class T>
void Squad<T>::addToBattlePredictor (
	BattlePredictor &battlePredictor
) const
{
	for (std::set<T*>::const_iterator i=m_units.begin() ; i!=m_units.end() ; i++) {
		battlePredictor.addUnit( (*i)->getBWAPIUnit() );
	}
}
